from __future__ import absolute_import
from collections import defaultdict, OrderedDict
import pdb
import numpy as np
from sklearn.metrics import average_precision_score
import os
import os.path as osp
from PIL import Image
from matplotlib import pyplot as plt

from ..misc import to_numpy

BASEDIR = osp.dirname(osp.abspath(__file__))


def _unique_sample(ids_dict, num):
    mask = np.zeros(num, dtype=np.bool)
    for _, indices in ids_dict.items():
        i = np.random.choice(indices)
        mask[i] = True
    return mask


def cmc(distmat,
        query_ids=None,
        gallery_ids=None,
        query_cams=None,
        gallery_cams=None,
        topk=100,
        separate_camera_set=False,
        single_gallery_shot=False,
        first_match_break=False):
    distmat = to_numpy(distmat)
    m, n = distmat.shape
    # Fill up default values
    if query_ids is None:
        query_ids = np.arange(m)
    if gallery_ids is None:
        gallery_ids = np.arange(n)
    if query_cams is None:
        query_cams = np.zeros(m).astype(np.int32)
    if gallery_cams is None:
        gallery_cams = np.ones(n).astype(np.int32)
    # Ensure numpy array
    query_ids = np.asarray(query_ids)
    gallery_ids = np.asarray(gallery_ids)
    query_cams = np.asarray(query_cams)
    gallery_cams = np.asarray(gallery_cams)
    # Sort and find correct matches
    indices = np.argsort(distmat, axis=1)
    matches = (gallery_ids[indices] == query_ids[:, np.newaxis])
    # Compute CMC for each query
    ret = np.zeros(topk)
    num_valid_queries = 0
    for i in range(m):
        # Filter out the same id and same camera
        valid = ((gallery_ids[indices[i]] != query_ids[i]) |
                 (gallery_cams[indices[i]] != query_cams[i]))
        if separate_camera_set:
            # Filter out samples from same camera
            valid &= (gallery_cams[indices[i]] != query_cams[i])
        if not np.any(matches[i, valid]):
            continue
        if single_gallery_shot:
            repeat = 10
            gids = gallery_ids[indices[i][valid]]
            inds = np.where(valid)[0]
            ids_dict = defaultdict(list)
            for j, x in zip(inds, gids):
                ids_dict[x].append(j)
        else:
            repeat = 1
        for _ in range(repeat):
            if single_gallery_shot:
                # Randomly choose one instance for each id
                sampled = (valid & _unique_sample(ids_dict, len(valid)))
                index = np.nonzero(matches[i, sampled])[0]
            else:
                index = np.nonzero(matches[i, valid])[0]
            delta = 1. / (len(index) * repeat)
            for j, k in enumerate(index):
                if k - j >= topk:
                    break
                if first_match_break:
                    ret[k - j] += 1
                    break
                ret[k - j] += delta
        num_valid_queries += 1
    if num_valid_queries == 0:
        raise RuntimeError("No valid query")
    return ret.cumsum() / num_valid_queries


def mean_ap(distmat,
            query_ids=None,
            gallery_ids=None,
            query_cams=None,
            gallery_cams=None):
    distmat = to_numpy(distmat)
    m, n = distmat.shape
    # Fill up default values
    if query_ids is None:
        query_ids = np.arange(m)
    if gallery_ids is None:
        gallery_ids = np.arange(n)
    if query_cams is None:
        query_cams = np.zeros(m).astype(np.int32)
    if gallery_cams is None:
        gallery_cams = np.ones(n).astype(np.int32)
    # Ensure numpy array
    query_ids = np.asarray(query_ids)
    gallery_ids = np.asarray(gallery_ids)
    query_cams = np.asarray(query_cams)
    gallery_cams = np.asarray(gallery_cams)
    # Sort and find correct matches
    indices = np.argsort(distmat, axis=1)
    matches = (gallery_ids[indices] == query_ids[:, np.newaxis])
    # Compute AP for each query
    aps = []
    for i in range(m):
        # Filter out the same id and same camera
        valid = ((gallery_ids[indices[i]] != query_ids[i]) |
                 (gallery_cams[indices[i]] != query_cams[i]))
        y_true = matches[i, valid]
        y_score = -distmat[i][indices[i]][valid]
        if not np.any(y_true):
            continue
        aps.append(average_precision_score(y_true, y_score))
    if len(aps) == 0:
        raise RuntimeError("No valid query")
    return np.mean(aps)


def read_image(path_img):
    if not osp.exists(path_img):
        raise IOError("The image path does not exist.")
    try:
        img = Image.open(path_img).convert("RGB")
    except:
        raise IOError("Something wrong with reading the image.")
    return img


def fetch_top_k(distmat,
                query_fpaths=None,
                gallery_fpaths=None,
                query_ids=None,
                gallery_ids=None,
                query_cams=None,
                gallery_cams=None,
                k=10):
    """
    Returns the indices of top k images in gallery that are most similar to the i_th query.
    Returns the items of top k images in gallery that are most similar to the i_th query.
    """
    sorted_indices = np.argsort(distmat, axis=1)  # ascending, small to big
    top_k_indices = sorted_indices[:, :k]

    q_drawer = []
    for i in range(top_k_indices.shape[0]):
        qi_topk_items = []
        for j in range(top_k_indices.shape[1]):
            fpath_ = gallery_fpaths[top_k_indices[i, j]]  # !!!
            id_ = gallery_ids[top_k_indices[i, j]]
            cam_ = gallery_cams[top_k_indices[i, j]]
            qi_topk_items.append({
                "path_img": fpath_,
                "id": id_,
                "camid": cam_
            })

        qi_result = {
            "q_self": {
                "path_img": query_fpaths[i],
                "id": query_ids[i],
                "camid": query_cams[i]
            },
            "match": qi_topk_items
        }
        q_drawer.append(qi_result)

    return top_k_indices, q_drawer


def show_searched_pics(q_drawer):
    fig_save_dir = osp.join(BASEDIR, 'saved_figures')
    if not osp.exists(fig_save_dir):
        os.makedirs(fig_save_dir)

    for i in range(len(q_drawer)):
        k = len(q_drawer[i]["match"])  # top k

        # 对于第i个query：
        qi_img = read_image(q_drawer[i]["q_self"]["path_img"])
        qi_id = q_drawer[i]["q_self"]["id"]
        qi_camid = q_drawer[i]["q_self"]["camid"]

        plt.figure(figsize=(20, 10))
        plt.subplot(1, k + 1, 1)
        plt.imshow(qi_img)
        plt.axis("off")
        # plt.title(f"query_id{qi_id}_cam{qi_camid}")
        plt.title(f"q{qi_id}")

        # 遍历它所匹配的检索结果
        for j in range(k):
            img = read_image(q_drawer[i]["match"][j]["path_img"])
            plt.subplot(1, k + 1, j + 2)
            plt.imshow(img)
            plt.axis("off")
            g_id = q_drawer[i]["match"][j]["id"]
            g_camid = q_drawer[i]["match"][j]["camid"]
            # plt.title(f"id{g_id}_cam{g_camid}")
            plt.title(f"{g_id}")

        # plt.tight_layout()
        # plt.show()
        save_path = osp.join(fig_save_dir, f"q{i}.png")
        plt.savefig(save_path)
        print(i, 'done')
        plt.close()
